#! /usr/bin/env python
# -*-coding: utf-8 -*-
#
# Date://08/17/2010 01:16:17 PM (CST)  
# Author: Gou Pengfei
#

import sys
import os
import optparse
import datetime
import re
from subprocess import call
from os.path import join as joinpath

progname = os.path.basename(sys.argv[0])

binary_path = os.path.join(os.environ['HOME'], 'ttools2-devel', 'bin')

int_benchmarks = ['gzip_graphic','vpr_route','gcc_166','mcf',
                  'perlbmk_makerand',
                  'crafty','parser_spec','gap','vortex_2',
                  'bzip2_program','twolf',]

float_benchmarks = ['wupwise','swim',
                    'mgrid','applu','mesa','art_110',
                    'equake','ammp','apsi'] 

qsub_conf_file_head = \
'''#! /usr/bin/env python

# AUTOMATICALLY GENERATED. DO NOT EDIT!

import sys
import os
import datetime
import re
from subprocess import call
from os.path import join as joinpath

var = '%s'
benchmark_to_run = '%s'
input_set = '%s'
conf = '%s'
script_name = '%s'
home = '%s'
'''

Usage = '''usage: %prog [options] binary name to run
Binaries to run are separated by single spaces.

Working cycle:

Step 1). Try '%prog -c' to compile M5
Step 2). Try '%prog hello' to let M5 simulator execute a 'hello' binary with default options'

Examples:

1) '%prog -i trips -m se -v fast -c': 
    Compile M5 with TRIPS ISA under SE mode and generate a fast version of M5.
2) '%prog -i trips -m se -v opt -t Exec hello': 
    Run M5 with TRIPS ISA under SE mode, using m5.opt and generating traces with trace flags 'Exec'

Tips:

1). Try '%prog -h' to check the default options.
2). Actually you can just run M5 command line directly without this script.
3). Please feel free to ask others for help.

Enjoy!
'''

lljob_template = '''
# @ executable = %s
# @ output = %s.out.$(jobid)
# @ error = %s.err.$(jobid)
# @ requirements = (OpSys == "Linux26")
# @ resources = ConsumableMemory(1500 mb)
# @ account_no = fsp2
# @ cpu_limit = 24:00:00
# @ class = 4hour
 
# @ queue
 '''

optparser = optparse.OptionParser(Usage)

optparser.add_option('-v', '--variants', dest='var',
                     default='opt',
                     help='Binary variant of m5 to run. This option can be fast, opt, or debug. '
                     " (default: '%default')" )
                     
optparser.add_option('-f', '--conf', dest='conf',
                     default='atomic-fast-mem',
                     help='Comma separated configurations of this run. Each configuration is repsented '
                     "by a script in ./configs/trips_config/ with a file name as '*.py'. "
                     "(default: '%default')")
                     
optparser.add_option('-p', '--confpath', dest='conf_path',
                     default='configs/trips_config',
                     help='Path of directory that holds configuration scripts.'
                     "(default: '%default')")
                     
optparser.add_option('-b', '--binarypath', dest='bin_path',
                     default=binary_path,
                     help='Path of the directory that holds all the binaries to run.'
                     " (default: '%default')" )

optparser.add_option('-d', '--outdir', dest='outdir',
                     default='./simout/trips',
                     help='Path of directory for output. '
                     " (default: '%default')" )

optparser.add_option('-t', '--traces', dest='traces',
                     default = '',
                     help='Comma separated trace flags of this run. Default value is left blank. ' )

optparser.add_option('-c', '--compile', action='store_true', dest='compile',
                     default = False,
                     help = 'Compile or not. ')        

optparser.add_option('-q', '--maui_queue', dest='queue',
                     default = 'workq',
                     help = 'Maui queue type.'
                     " (default queue: '%default')")        

optparser.add_option('--batch', action='store_true',
                     dest='batch',
                     default = False,
                     help = 'Running spec using batch job'
                     " (default: '%default')")        

optparser.add_option('--script_only', action='store_true',
                     dest='script_only',
                     default = False,
                     help = 'Generating spec script only, do not run'
                     'them.')        

optparser.add_option('-i', '--inputset', dest='input_set',
                     default = 'ref',
                     help='Input set used by spec running'
                     " (default: '%default')")

optparser.add_option('-C', '--cpu_models', dest='cpu_models',
                     default =\
                     'SimpleEdgeCPU,AtomicEdgeCPU,DistributedEdgeCPU',
                     help = '''The options for comma separated cpu
                     models.''' "(default: '%default')")
                     
optparser.add_option('--compile_options', dest='c_options',
                     default = '',
                     help = 'The options for compiling M5 '
                     "(default: '%default')")  
                     
optparser.add_option('-e', '--extra', dest='extras',
                     default = '',
                     help='Extra execution options. Default value is left'
                     "blank. ")

optparser.add_option('-I', '--isa', dest='isa',
                     default = 'TRIPS',
                     help='The ISA of M5 to run. '
                     "(default: '%default')")       

optparser.add_option('-m', '--mode', dest='mode',
                     default = 'SE',
                     help='The mode of M5 to run. '
                     "(default: '%default')")   

optparser.add_option('-l', '--lljob', dest='lljob',
                     default='job.cmd',
                     help='Name of output lljob script'
                     " (default: '%default')" )
 
(options,leftovers) = optparser.parse_args()

compile_template = "scons %s build/%s_%s/m5.%s"

cmd_template = "build/%s_%s/m5.%s --outdir=%s --stats-file=\
stats.txt"

binary_cmd_template = " %s/%s.py --simplebench %s"

to_run = leftovers

# Call os.system() and raise exception if return status is non-zero
def system(cmd):
    try:
        retcode = call(cmd, shell=True)
        if retcode < 0:
            print >>sys.stderr, "Child was terminated by signal", -retcode
            print >>sys.stderr, "When attemping to execute: %s" % cmd
            sys.exit(1)
        elif retcode > 0:
            print >>sys.stderr, "Child returned", retcode
            print >>sys.stderr, "When attemping to execute: %s" % cmd
            sys.exit(1)
    except OSError, e:
        print >>sys.stderr, "Execution failed:", e
        print >>sys.stderr, "When attemping to execute: %s" % cmd
        sys.exit(1)

class QsubConfFile(object):
    def __init__(self, var, bench_name, input_set, conf_name,
            home_path):
        self.var = var
        self.input_set = input_set
        self.conf_name = conf_name
        self.bench_name = bench_name
        self.home_path = home_path
        self.script_path = ""

        sys.path.append(joinpath(home_path, 'M5-EDGE', 'configs',
                                 'common'))
        from cpu2000 import simout_dir
        self.output_dir = joinpath(simout_dir, options.isa.lower(), 'linux', 
                                   input_set, (bench_name.split('_'))[0],
                                   ((self.conf_name).split('/'))[-1])

        if not os.path.isdir(self.output_dir):
            os.makedirs(self.output_dir)

        # Remove the flag files
        if os.path.isfile(joinpath(self.output_dir, '.error')):
            os.remove(joinpath(self.output_dir, '.error'))

        if os.path.isfile(joinpath(self.output_dir, '.running')):
            os.remove(joinpath(self.output_dir, '.running'))

        if os.path.isfile(joinpath(self.output_dir, 'stats.txt')):
            os.remove(joinpath(self.output_dir, 'stats.txt'))

    def generate(self):
        head = qsub_conf_file_head % (self.var,
                                      self.bench_name,
                                      self.input_set,
                                      ((self.conf_name).split('/'))[-1],
                                      self.conf_name,
                                      self.home_path)

        home = os.environ['HOME']
        m5_dir = joinpath(home, 'M5-EDGE')
        script_dir = joinpath(m5_dir, 'util', 'qsub', ((self.conf_name).split('/'))[-1])
        self.script_path = str(script_dir)

        if not os.path.isdir(script_dir):
            os.makedirs(script_dir)

        output_file = open(joinpath(script_dir, self.bench_name + '.py'), 'w')
        output_file.write(head)

        template = open(joinpath(home, 'M5-EDGE', 'util', 'qsub', 'qsub-template'), 'r')
        output_file.write(template.read())
        output_file.close()

    def scriptpath(self):
        return self.script_path

if __name__ == '__main__':
    isa = options.isa
    mode = options.mode

    isa = isa.upper()
    mode = mode.upper()

    if options.batch:

        cmd_template = 'qsub -q %s -o %s -e %s %s/util/qsub/%s/%s.py &'

        home = os.environ['HOME']
        m5_dir = joinpath(home, 'M5-EDGE')

        confs = (options.conf).split(',')

        if isa != 'TRIPS' and mode != 'SE':
            print '''We only support TRIPS_SE mode for running
                  spec using qsub'''
            sys.exit(1)

        if 'all' in to_run:
            benches = int_benchmarks + float_benchmarks
        elif 'int' in to_run:
            benches = int_benchmarks
        elif 'fp' in to_run:
            benches = float_benchmarks
        else:
            benches = to_run

        # if this configuration is a directory in
        # $HOME/M5_EDGE/configs/trips_config, then all the files in
        # this directory will be executed as configurations
        conf_dir = joinpath(m5_dir, 'configs', 'trips_config')
        configurations = []
        for conf in confs:
            if os.path.isdir(joinpath(conf_dir, conf)):
                all_files = os.listdir(joinpath(conf_dir, conf))
                all_files = [joinpath(conf, i) for i in all_files]
                configurations.extend(all_files)
            else:
                configurations.extend([conf])
        
        configurations.sort()

        lljob = None
        if options.lljob != "":
            lljob = open(options.lljob, 'w')

        for conf in configurations: 
            conf = (conf.split('.'))[0]
            for bench in benches:
                conf_file = QsubConfFile(options.var,
                                         bench,
                                         options.input_set,
                                         conf,
                                         home)
                conf_file.generate()

                qsub_output = joinpath(conf_file.output_dir, 'qsub.out')
                qsub_errout = joinpath(conf_file.output_dir, 'qsub.err')

                cmd = cmd_template % (options.queue,
                                      qsub_output,
                                      qsub_errout,
                                      m5_dir,
                                      (conf.split('/'))[-1],
                                      bench)
                print cmd

                job_name = bench + '-' + conf.split('/')[-1]
                job = lljob_template % (conf_file.scriptpath() + '/' + bench + '.py', job_name, job_name)

                print >>lljob, job

                if not options.script_only:
                    try:
                        system(cmd)
                    except:
                        print >>sys.stderr, "Incorrect command line for M5."
                        sys.exit(1)

    else:
        if options.compile:
            cpu_models = 'CPU_MODELS=' + options.cpu_models
            compile_options = cpu_models + ' ' + options.c_options
            compile_cmd = compile_template % (compile_options, isa, mode, options.var)

            print "Running command"
            print "Command: %s" % compile_cmd
            print ">>>>>>>>"
            
            try:
                system(compile_cmd)
            except:
                print >>sys.stderr, "Incorrect compilation command for M5"
                sys.exit(1)
                
            sys.exit(0)

        for bin in to_run:
        
            cmd = cmd_template
            
            if options.traces != '':
                cmd += (' --trace-file=trace.s --trace-flags=%s' % options.traces)

            cmd = cmd % (isa, mode, options.var, options.outdir)
            binary_cmd = binary_cmd_template % (options.conf_path, options.conf, bin)

            if options.extras != '':
                cmd += ' ' + options.extras

            cmd += binary_cmd

            print "Running command"
            print "Command: %s" % cmd
            print "<<<<<<<<"

            try:
                system(cmd)
            except:
                print >>sys.stderr, "Incorrect command line for M5."
                sys.exit(1)

sys.exit(0)

